# -*- coding: utf-8 -*-
"""
tyens:
    lon
    lat
    umax
    pmin
    time
tyobs:
    lon
    lat
    umax
    pmin 
    time
"""

import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall
from geopy.distance import geodesic
import matplotlib.ticker as ticker
import seaborn as sns
import xarray as xr

ini_time_mode = ['00','12']
# names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
names = ['dusurui_16']
tynames,tyids = split_str_id(names)
draw_obs_opt = True
obs_baseline='land'  # ‘land’ 'RI'
RIstd = 7
tyrmse = {}
show = False


track_id = np.arange(1,52)

# 遍历台风种类
for ty,tyid in zip(tynames,tyids):
    ''' init '''
    dt = timedelta(days=1)
    directory =  f'D:\\met_data\\ty_ensemble\\{ty}_{tyid}\\'    # 要筛选的文件开头
    dates_name = os.listdir(directory)
    dates_name = [i for i in dates_name if i[-2:] in ini_time_mode ]
    obs_path = rf'D:\met_data\ty_obs\{ty}_CMAobs.txt'
    tyobs = tydat_CMA(obs_path)
    # 准备数组 xr.DataArray
    da = xr.DataArray(np.nan,dims=['variable','start_time','track_id'],coords={'variable':['umax','pmin'],'start_time':dates_name,'track_id':track_id} )
    # 遍历起报时间
    for date_name in tqdm(dates_name,total=len(dates_name)):
        dir_date = os.path.join(directory,date_name)
        name_ensembles = os.listdir(dir_date)
        path_ensembles = [os.path.join(dir_date,i) for i in name_ensembles if i.startswith('TRACK')]
        # 遍历某台风某起报时间的所有集合预报成员
        for path in path_ensembles:
            t_id = int(path.split('_')[-1])
            if t_id==0:
                continue # ERA5 情况
            else:
                # 算RMSE
                tyens = tydat(path,RIstd)
                # 求时间交集
                intersect_time = np.intersect1d(tyobs.time, tyens.time)
                bool_obs,bool_ens = np.isin(tyobs.time,intersect_time),np.isin(tyens.time,intersect_time)
                umax_obs, umax_ens = tyobs.umax[bool_obs],tyens.umax[bool_ens]
                pmin_obs, pmin_ens = tyobs.pmin[bool_obs],tyens.pmin[bool_ens]
                lat_obs , lat_ens  = tyobs.lat[bool_obs],tyens.lat[bool_ens]
                lon_obs , lon_ens  = tyobs.lon[bool_obs],tyens.lon[bool_ens]
                da.loc['umax',date_name,t_id] = np.sqrt(np.mean((umax_obs-umax_ens)**2))
                da.loc['pmin',date_name,t_id] = np.sqrt(np.mean((pmin_obs-pmin_ens)**2))
    tyrmse[ty] = da
            







# 对起报时间求平均的图
for ty,tyid in zip(tynames,tyids):
    for var in ['umax','pmin']:
        a = tyrmse[ty].loc[var].values
        a = np.where(a>10000,np.nan,a)
        a = np.ma.masked_invalid(a)
        # 创建图形和坐标轴
        fig, ax = plt.subplots(figsize=(20, 10))
        # 绘制美化的箱线图
        box = ax.boxplot(
            a,
            notch=True,                 # 中位数处切口
            patch_artist=True,          # 填充盒子
            widths=0.6,                 # 箱体宽度
            boxprops=dict(facecolor='lightblue', edgecolor='navy', linewidth=1.5),
            medianprops=dict(color='red', linewidth=2),
            whiskerprops=dict(color='navy', linewidth=1.2),
            capprops=dict(color='navy', linewidth=1.2),
            flierprops=dict(marker='o', markerfacecolor='orange',
                            markeredgecolor='gray', markersize=5, linestyle='none')
        )
        # 设置标题和轴标签
        ax.set_title(f'{ty}  mean {var}-rmse over ensembles', fontsize=20, fontweight='bold')
        ylabel = f'{var}/hPa' if var=='pmin' else f'{var}/m/s'
        ax.set_ylabel(ylabel, fontsize=16,fontweight='bold') 
        ax.set_xlabel('Ensemble member',fontsize=16,fontweight='bold')

        ax.yaxis.grid(True, linestyle='--', alpha=0.7)
        plt.tight_layout()
        if show:
            plt.show()
        else:
            plt.savefig(rf'C:\Users\lenovo\Desktop\typic\RMSE\box_{ty}_{var}.png',dpi=600)
            plt.close(fig)

    
    

    


# 变量对应的颜色和标记
var_props = {
    'umax': dict(color='#1f77b4', marker='o'),
    'pmin': dict(color='#ff7f0e', marker='s'),
}

for ty, tyid in zip(tynames, tyids):
    # 新建一张图
    fig, ax = plt.subplots(figsize=(10, 6))
    for var, props in var_props.items():
        # 读取数组并屏蔽超大值和 NaN
        a = tyrmse[ty].loc[var].values
        a = np.where(a > 10000, np.nan, a)
        a = np.ma.masked_invalid(a).T
        
        # 计算每个“起报时间”对应的平均值
        y = np.mean(a, axis = 0)
        x = np.arange(len(y))  # 如果你有具体的时间标签，可以改成真实的 time_index
        
        # 绘制折线
        ax.plot(
            x, y,
            label=var,
            linewidth=2,
            marker=props['marker'],
            markersize=6,
            color=props['color']
        )
    
    # 设置标题、标签、图例
    ax.set_title(f'{ty} — Umax / Pmin Ensemble_mean', fontsize=16, fontweight='bold')
    ax.set_xlabel('Forecaset start time', fontsize=14)
    ax.set_ylabel('Ensemble mean rmse', fontsize=14)
    ticks = np.arange(len(tyrmse[ty].start_time))
    labels = tyrmse[ty].start_time.values

    ax.set_xticks(ticks)                # 先告诉 Matplotlib 有多少个刻度
    ax.set_xticklabels(labels, rotation=60, ha='right',fontsize=8)
    ax.legend(title='variable', fontsize=12, title_fontsize=12, loc='best')
    ax.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    if show:
        plt.show()
    else:
        plt.savefig(rf'C:\Users\lenovo\Desktop\typic\RMSE\plot_{ty}_{var}.png',dpi=600)
        plt.close(fig)
    
    
    
    


for ty, tyid in zip(tynames, tyids):
    for var in ['umax', 'pmin']:
        # 1) 读取数据
        da = tyrmse[ty].loc[var]                     # DataArray (track_id, start_time)
        data = da.values.astype(float)               # 转 float，避免 int 溢出
        data = np.where(data > 10000, np.nan, data)  # 屏蔽异常值
        data = np.ma.masked_invalid(data)            # NaN → masked

        # 2) 坐标
        members = da.track_id.values
        times   = da.start_time.values

        # 3) 画图
        fig, ax = plt.subplots(figsize=(12, 8))
        im = ax.imshow(
            data,                      # 已经是 (track_id, start_time)
            aspect='auto',
            interpolation='none',
            cmap='viridis',
            origin='upper'             # 与坐标顺序一致
        )

        # ⭐ 刻度与标签一一对应
        ax.set_xticks(np.arange(len(members)))
        ax.set_xticklabels(members, rotation=0, fontsize=8)
        ax.set_yticks(np.arange(len(times)))
        ax.set_yticklabels(times,fontsize=8)
                           

        # ⭐ 标题 & 标签
        ax.set_title(f'{ty} — {var} RMSE Heatmap', fontsize=16, fontweight='bold')
        ax.set_xlabel('Ensemble Member', fontsize=14)
        ax.set_ylabel('Forecast Start Time', fontsize=14)

        # ⭐ 色标
        cbar = fig.colorbar(im, ax=ax)
        cbar.set_label('RMSE', fontsize=12)

        plt.tight_layout()
        if show:
            plt.show()
        else:
            plt.savefig(
                rf'C:\Users\lenovo\Desktop\typic\RMSE\heatmap_{ty}_{var}.png',
                dpi=600
            )
            plt.close()